from .dataset_utils import *
from enum import Enum


class TokenColor(Enum):
    NOCOLOR = 0
    BLUE = 1
    GREEN = 2
    RED = 3


class Token(str):
    """
    class继承自str, str所有方法都可以直接拿过来用
    """

    def __init__(self, token_str, color=TokenColor.NOCOLOR):
        super(Token, self).__init__()
        self.color = color  # for Agile Syntax
        self.du = None  # token is DEF or USE
        self.loc = None  # token location
        self.local = False  # global or local means if token in a scope
        self.to = None  # data flow from here to one location


class Record(list):
    """

    """

    def __init__(self, token_strs):
        # token_str是否是Token类的对象，如果不是，将它转成Token类型的
        lines = [token_str if isinstance(token_str, Token) else Token(token_str) for token_str in token_strs]
        super(Record, self).__init__(lines)


def to_records(record_str_list):
    return [Record(record_str) for record_str in record_str_list]


class Location(object):
    """
    记录一个token当前出现的位置信息：
    文件名
    哪个记录
    第几个token
    """
    instance = {}
    check_ = False

    @staticmethod
    def create_location(filename, record_id, token_id):
        key = (filename, record_id, token_id)
        if key not in Location.instance:
            Location.check_ = True
            Location.instance[key] = Location(filename, record_id, token_id)
            Location.check_ = False
        return Location.instance[key]

    def __init__(self, filename, record_id, token_id):
        assert Location.check_, "constructor must be involved by Location.create_location(...)"
        self.filename = filename
        self.record_id = record_id
        self.token_id = token_id

    def __str__(self):
        return f"{self.filename}:{self.record_id}：{self.token_id}"  # 类似c++中的占位符: %s %d

    #         "{}/hello{}".format(100, 100.2)
    #         a=100
    #         b='xyzxyz'
    #         "{}/hello{}".format(a,b)
    #         f"{a}/hello{b}"

    def __repr__(self):
        return self.__str__()


class Dataset(object):
    def __init__(self, dir,
                 annotation_start="//",
                 record_start_keywords=["class", "def", "let", "defm",
                                        "multiclass", "foreach"],
                 record_one_line_keywords=["include"]
                 ):
        # 每个文件与其记录做一个映射，形成一个字典； 或者文件名与其记录形成一个tuple，所有文件构成一个list
        self.documents_dict = Dataset.get_records(dir, annotation_start,
                                                  record_start_keywords, record_one_line_keywords)
        self.documents_list = list(self.documents_dict.items())

    @staticmethod
    def get_records(dir,
                    annotation_start="//",
                    record_start_keywords=["class", "def", "let", "defm",
                                           "multiclass", "foreach"],
                    record_one_line_keywords=["include"]
                    ):
        """
        静态函数，由staticmethod定义
        加上该限制，类Dataset不需要实例化便可以直接调用该函数
        """
        documents = {}
        if os.path.isdir(dir):
            for home, dirs, files in os.walk(dir):
                for file in files:
                    fullname = file  # os.path.join(home, file)
                    if fullname.endswith('.td'):
                        record_list = get_document_records(fullname, annotation_start,
                                                           record_start_keywords, record_one_line_keywords)
                        documents[fullname] = to_records(record_list)  # 文件名和所有record的映射
        else:
            # 传入的不是目录，直接是文件名
            record_list = get_document_records(dir)
            documents[dir] = to_records(record_list)
        return documents

    def __getitem__(self, idx):
        if isinstance(idx, int):
            return self.documents_list[idx]
        elif isinstance(idx, str):
            return self.documents_dict[idx]
        elif isinstance(idx, Location):
            return self.documents_dict[idx.filename][idx.record_id]
        else:
            raise ValueError("only support int or str as index")  # 手动引发异常

    def get_document(self, loc: Location):
        return self.documents_dict[loc.filename]

    def get_record(self, loc: Location):
        return self.documents_dict[loc.filename][loc.record_id]
